import pickle

import torch
from matplotlib import pyplot as plt, ticker

from problems.problem import Problem
from utils.file import get_root_path
from websocket.msghandler import MsgContext, MessageType
from websocket.msghandler.multipleExcelHandler import MultipleExcelHandler


class MultipleDrawingHandler:
    pdf_file_path = get_root_path() + r'\pdf'

    @classmethod
    def handle(cls, message):
        draw_type = message['type']
        selected_algorithms = message['selectedAlgorithms']
        selected_problems = message['selectedProblems']
        algorithm_size = len(selected_algorithms)
        problem_size = len(selected_problems)
        indicator_data = message['indicator']
        indicator_values = MsgContext.multiple_indicators[indicator_data['name']]
        populations = MsgContext.multiple_populations
        problem_objs = MsgContext.multiple_problem_objs
        var_dim = message['var_dim']
        obj_dim = message['obj_dim']
        run_size = message['run_size']
        save_pdf = message['savePdf']
        # 创建图表和子图网格
        fig_size = (4 * algorithm_size, 3 * problem_size)
        if draw_type == 'obj' or draw_type == 'dec':
            if (draw_type == 'obj' and obj_dim != 3) or (draw_type == 'dec' and var_dim != 3):
                fig, axs = plt.subplots(problem_size, algorithm_size,
                                        figsize=fig_size,
                                        squeeze=False)  # figsize可以根据需要调整大小
            else:
                fig, axs = plt.subplots(problem_size, algorithm_size,
                                        figsize=fig_size,
                                        subplot_kw={'projection': '3d'},
                                        squeeze=False)  # figsize可以根据需要调整大小
            for row_index, problem in enumerate(selected_problems):
                problem_index = problem['key']
                problem_obj: Problem = problem_objs[problem_index]
                for column_index, algorithm in enumerate(selected_algorithms):
                    algorithm_index = algorithm['key']
                    indicator_value = indicator_values[problem_index][algorithm_index]
                    _, best_index = torch.median(indicator_value, 0)
                    population = populations[problem_index][algorithm_index][best_index][-1]
                    ax = axs[row_index, column_index]
                    if (draw_type == 'obj' and obj_dim == 2) or (draw_type == 'dec' and var_dim == 2):
                        ax.scatter(population[draw_type][:, 0].tolist(),
                                   population[draw_type][:, 1].tolist(),
                                   color='red',
                                   label='Your'
                                   )
                        ax.scatter(problem_obj.optimal_solutions[:, 0].tolist(),
                                   problem_obj.optimal_solutions[:, 1].tolist(),
                                   color='blue',
                                   label='PF'
                                   )
                        ax.set_xlabel('f1')
                        ax.set_ylabel('f2')
                        ax.legend()
                    elif (draw_type == 'obj' and obj_dim == 3) or (draw_type == 'dec' and var_dim == 3):
                        ax.scatter(population[draw_type][:, 0].tolist(),
                                   population[draw_type][:, 1].tolist(),
                                   population[draw_type][:, 2].tolist(),
                                   color='red',
                                   label='Your'
                                   )
                        ax.scatter(problem_obj.optimal_solutions[:, 0].tolist(),
                                   problem_obj.optimal_solutions[:, 1].tolist(),
                                   problem_obj.optimal_solutions[:, 2].tolist(),
                                   color='blue',
                                   label='PF'
                                   )
                        ax.set_xlabel('f1')
                        ax.set_ylabel('f2')
                        ax.set_zlabel('f3')
                        ax.legend()
                    else:
                        for index in range(len(population[draw_type])):
                            ax.plot(population[draw_type][index], color='blue')
                            ax.set_xlabel('column index')
                            ax.set_ylabel('value')
                    ax.set_title(
                        f'{algorithm["label"]} on {problem["label"]}')  # 给每个子图设置标题
        else:
            fig, axs = plt.subplots(problem_size, algorithm_size,
                                    figsize=fig_size,
                                    squeeze=False)  # figsize可以根据需要调整大小
            indicator_values_list = MsgContext.indicator_value_list.get(indicator_data['name'],
                                                                        MsgContext.create_2D_list(
                                                                            MsgContext.tensor_size[0],
                                                                            MsgContext.tensor_size[1]))
            for row_index, problem in enumerate(selected_problems):
                problem_index = problem['key']
                for column_index, algorithm in enumerate(selected_algorithms):
                    algorithm_index = algorithm['key']
                    all_run_value = []
                    ax = axs[row_index][column_index]
                    if indicator_values_list[problem_index][algorithm_index] is None:
                        for nth_run in range(run_size):
                            file_path = MsgContext.file_path_list[problem_index][algorithm_index][nth_run]
                            with open(file_path, 'rb') as f:
                                load_data = pickle.load(f)
                            populations = load_data[0]
                            temp_indicator_value = []
                            for population in populations:
                                indicator_class = MsgContext.get_cls_by_str(indicator_data['cls'])
                                indicator_value = indicator_class.get_value(population,
                                                                            MsgContext.multiple_problem_objs[
                                                                                problem_index])
                                temp_indicator_value.append(indicator_value)
                            all_run_value.append(temp_indicator_value)
                        indicator_values_list[problem_index][algorithm_index] = torch.tensor(all_run_value).mean(
                            dim=0)
                    ax.plot(MsgContext.multiple_FE[problem_index][algorithm_index],
                            indicator_values_list[problem_index][algorithm_index],
                            color='blue',
                            )
                    ax.set_title(f'{algorithm["label"]} on {problem["label"]}')  # 给每个子图设置标题
                    ax.set_xlabel('evaluation')
                    ax.set_ylabel('value')
                    # 设置 x 轴刻度为科学计数法
                    formatter = ticker.ScalarFormatter(useMathText=True)
                    formatter.set_scientific(True)
                    formatter.set_powerlimits((-1, 2))
                    ax.xaxis.set_major_formatter(formatter)
            MsgContext.indicator_value_list[indicator_data['name']] = indicator_values_list
        # 设置整体图形的标题
        # fig.suptitle(indicator_data['name'])
        plt.tight_layout()
        plt.show()
        if save_pdf:
            file_abs_path = MultipleExcelHandler.get_file_path(cls.pdf_file_path, suffix='.pdf')
            fig.savefig(file_abs_path)
            MsgContext.task_complete_return(MessageType.MULTI_Drawing_COMPLETE, f"file is saved in {file_abs_path}")
            return
        MsgContext.task_complete_return(MessageType.MULTI_Drawing_COMPLETE, f"table has been generated")
